package com.yc.config;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.transport.rest_client.RestClientTransport;
import com.yc.services.ToolServices;
import dev.langchain4j.community.model.dashscope.QwenEmbeddingModel;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.mcp.McpToolProvider;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.*;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.elasticsearch.ElasticsearchEmbeddingStore;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import org.elasticsearch.client.RestClient;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.bind.annotation.RequestParam;

@Configuration
public class AiConfig {

    // 定义AI服务接口（注意保持命名一致）
    public interface AiAssistant {

        String chat(@MemoryId int memoryId, @UserMessage String question);

        TokenStream chatStream(@MemoryId int memoryId, @UserMessage String question);

        TokenStream chatdeepseekStream(@MemoryId String memoryId, @UserMessage String question);
    }

//    @Bean
//    public InMemoryEmbeddingStore<TextSegment> embeddingStore() {
//        return new InMemoryEmbeddingStore<>();
//    }
    @Bean
    public EmbeddingStore<TextSegment> elasticsearchEmbeddingStore(ElasticsearchClient elasticsearchClient,
                                                               QwenEmbeddingModel qwenEmbeddingModel) {
        // 从 ElasticsearchClient 中获取底层的 RestClient
       RestClient restClient = ((RestClientTransport) elasticsearchClient._transport()).restClient();

        /**
         * 构建并返回一个基于Elasticsearch的向量嵌入存储实例，
         * 用于存储文本嵌入向量并支持向量相似度搜索。
         */
        return ElasticsearchEmbeddingStore.builder()
                // 配置Elasticsearch REST客户端，用于与Elasticsearch集群通信
                // restClient应已预先配置好集群地址、认证信息等
                .restClient(restClient)
                // 指定存储向量数据的Elasticsearch索引名称
                // 索引类似于数据库中的表，用于组织和存储文档
                .indexName("langchain4j-embeddings")
                // 设置向量维度，必须与嵌入模型生成的向量维度一致
                // qwenEmbeddingModel.dimension()返回Qwen模型生成的向量维度（如768、1536等）
                .dimension(qwenEmbeddingModel.dimension())
                // 构建并初始化ElasticsearchEmbeddingStore实例
                .build();
    }

    @Bean
    public AiAssistant aiAssistant(ChatModel chatModel,
                                   StreamingChatModel streamingChatModel,
                                   ChatMemoryStore chatMemoryStore,
                                   ToolServices tools,
                                   EmbeddingStore embeddingStore,
                                   QwenEmbeddingModel qwenEmbeddingModel,
                                   McpToolProvider mcpToolProvider) {

        ChatMemoryProvider chatMemoryProvider = memoryId -> MessageWindowChatMemory.builder()
                .id(memoryId)
                .maxMessages(1000)
                .chatMemoryStore(chatMemoryStore)
                .build();

        EmbeddingStoreContentRetriever retriever = EmbeddingStoreContentRetriever.builder()
                .embeddingModel(qwenEmbeddingModel)
                .embeddingStore(embeddingStore)
                .build();

        return AiServices.builder(AiAssistant.class)
                .chatModel(chatModel).streamingChatModel(streamingChatModel)
                .chatMemoryProvider(chatMemoryProvider)
                .tools(tools)
                .toolProvider(mcpToolProvider) // 添加MCP工具提供者
                .contentRetriever(retriever)
                .build();
    }
}
